import os
import argparse
import sys

def find_txt_with_specific_lines(directory):
    result = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(".txt"):
                file_path = os.path.join(root, file)
                try:
                    with open(file_path, "r", encoding="utf-8") as f:
                        # 过滤掉空行并统计
                        lines = [line.strip() for line in f if line.strip()]
                        if len(lines) in (0, 1, 2, 6):
                            result.append(file_path)
                except Exception as e:
                    print(f"读取文件失败: {file_path}, 错误: {e}")
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="查找指定目录下行数为 0,1,2,6 的 txt 文件（忽略空行）")
    parser.add_argument("--path", type=str, required=True, help="需要扫描的目录路径")
    args = parser.parse_args()

    directory = args.path

    # 检查目录是否存在
    if not os.path.isdir(directory):
        print(f"错误: 输入的目录不存在: {directory}")
        sys.exit(1)

    files = find_txt_with_specific_lines(directory)
    if files:
        print("符合条件的文件路径:")
        for f in files:
            print(f)
    else:
        print("没有找到符合条件的文件。")
